#include "headfiles.h"

int trap(int a[], int n) {
    if (n <= 1) {
        return 0;
    }
    
    // peak height and location
    vector<int> height;
    vector<int> pos;
    bool first_border = false;
    int incr;
    
    for (int i = 1; i < n; i++) {
        if (first_border == false) {
            if (a[i] < a[i - 1]) {
                height.push_back(a[i - 1]);
                pos.push_back(i - 1);
                first_border = true;
                incr = false;
            }
        }
        else {
            if (incr) {
                if (a[i] < a[i - 1]) {
                    height.push_back(a[i - 1]);
                    pos.push_back(i - 1);
                    incr = false;
                }
            }
            else {
                if (a[i] > a[i - 1]) {
                    height.push_back(a[i - 1]);
                    pos.push_back(i - 1);
                    incr = true;
                }
            }
        }
    }
    
    if (incr) {
        height.push_back(a[n - 1]);
        pos.push_back(n - 1);
    }
    
    int size = pos.size();
    int left_height = height[0];
    int left = pos[0];
    int right_height, right;
    int water = 0;
    for (int i = 1; i < size; i++) {
        if (left_height < height[i]) {
            right = pos[i];
            left++;
            while (left < right) {
                water += left_height - a[left];
                left++;
            }
            left_height = height[i];
            left = right;
        }
        else { // Find the highest right side
            right_height = -1;
            int j;
            for (j = i; j < size; j++) {
                if (right_height < height[j]) {
                    right_height = height[j];
                    right = pos[j];
                    if (right_height >= left_height) {
                        break;
                    }
                }
            }
            
            int height = (left_height > right_height) ? right_height : left_height;
            left++;
            while (left < right) {
                int tmp_water = height - a[left];
                tmp_water = (tmp_water > 0) ? tmp_water : 0;
                water += tmp_water;
                left++;
            }
            left_height = right_height;
            left = right;
            i = right;
        }
    }
    
    return water;
}

